#!/usr/bin/env python
# coding: utf-8

# In[ ]:


import operator
import struct
import numpy as np
import heapq
import random
train_images_dir = 'train-images.idx3-ubyte'# 训练集
train_labels_dir = 'train-labels.idx1-ubyte'# 训练集标签
test_images_dir = 't10k-images.idx3-ubyte'# 测试集
test_labels_dir = 't10k-labels.idx1-ubyte'# 测试集标签


# In[ ]:


def decode_idx3_ubyte(idx3_ubyte_file):
    bin_data = open(idx3_ubyte_file, 'rb').read()
    # 解析文件头信息，依次为魔数、图片数量、每张图片高、每张图片宽
    offset = 0
    fmt_header = '>iiii'
    magic_number, num_images, num_rows, num_cols = struct.unpack_from(fmt_header, bin_data, offset)
    print('魔数:%d, 图片数量: %d张, 图片大小: %d*%d' % (magic_number, num_images, num_rows, num_cols))
    # 解析数据集
    image_size = num_rows * num_cols
    offset += struct.calcsize(fmt_header)
    fmt_image = '>' + str(image_size) + 'B'
    images = np.zeros([num_images,num_rows,num_cols])
    for i in range(num_images):
        images[i] = np.array(struct.unpack_from(fmt_image, bin_data, offset)).reshape((num_rows,num_cols))
        offset += struct.calcsize(fmt_image)
    return images
 
def decode_idx1_ubyte(idx1_ubyte_file):
    bin_data = open(idx1_ubyte_file, 'rb').read()
    #解析文件头信息，依次为魔数和标签数
    offset = 0
    fmt_header = '>ii'
    magic_number, num_images = struct.unpack_from(fmt_header, bin_data, offset)
    print('魔数:%d, 图片数量: %d张' % (magic_number, num_images)) #解析数据集
    offset += struct.calcsize(fmt_header)
    fmt_image = '>B'
    labels = np.zeros(num_images)
    for i in range(num_images):
        labels[i]=struct.unpack_from(fmt_image, bin_data, offset)[0]
        offset += struct.calcsize(fmt_image)
    return labels


# In[ ]:


def average_pooling(feature_map,size,stride):
    channel=feature_map.shape[0]
    height=feature_map.shape[1]
    width=feature_map.shape[2]
    padding_height=np.uint16(round((height-size+1)/stride))
    padding_width=np.uint16(round((width-size+1)/stride))
    print(padding_height,padding_width)
    pool_out = np.zeros((channel,padding_height,padding_width))
    for map_num in range(channel):  
        out_height = 0  
        for r in np.arange(0,height, stride):  
            out_width = 0  
            for c in np.arange(0, width, stride):  
                pool_out[map_num,out_height, out_width] = np.sum(feature_map[map_num,r:r+size,c:c+size])/4
                out_width=out_width+1
            out_height=out_height+1
    return pool_out

def max_pooling(feature_map,stride):
    channel=feature_map.shape[0]
    height=feature_map.shape[1]
    width=feature_map.shape[2]
    pool_out = np.zeros((channel,height,width//2))
    for map_num in range(channel): 
        out_height = 0  
        for r in np.arange(0,height):  
            out_width = 0  
            for c in np.arange(0, width, stride):  
                pool_out[map_num,out_height, out_width] = np.max(feature_map[map_num,r,c:c+stride])
                out_width=out_width+1
            out_height=out_height+1
    return pool_out


# In[ ]:


def classify(X,dataset,labels,k):
    datasetsize = dataset.shape[0]#以下距离计算公式
    distances_2 = np.sum(np.array(np.tile(X,(datasetsize,1))-dataset)**2,1)
    minheap = []        #基于堆排序的Top-K，不对全部的距离排序
    for i in range(len(distances_2)):
        val=distances_2[i]
        if len(minheap) >= k:
            min_val = minheap[0][0]
            if -val < min_val :  #取负数之后，最大的相当于原来的最小
                pass
            elif val>0.0:
                heapq.heapreplace(minheap,[-val,i])
        elif val>0.0:                  #去掉距离为0的邻居，即自己
            heapq.heappush(minheap, [-val,i])
    #print(minheap)
    classCount = {} #前K个距离最小的
    for vec in minheap:
        voteIlabel = labels[vec[1]]
        classCount[voteIlabel] = classCount.get(voteIlabel, 0) + 1
    sortedClassCount = sorted(classCount.items(), key=operator.itemgetter(1), reverse=True)
    return sortedClassCount[0][0]


# In[ ]:


train_images = decode_idx3_ubyte(train_images_dir)
train_labels = decode_idx1_ubyte(train_labels_dir)
print(train_images.shape)
print(train_labels.shape)
test_images = decode_idx3_ubyte(test_images_dir)
test_labels = decode_idx1_ubyte(test_labels_dir)
print(test_images.shape)
print(test_labels.shape)


# In[ ]:


train_images=average_pooling(train_images,2,2)
train_images=max_pooling(train_images,2)
train_images=train_images.reshape(60000,-1)
print(train_images.shape)
test_images=average_pooling(test_images,2,2)
test_images=max_pooling(test_images,2)
test_images=test_images.reshape(10000,-1)
print(test_images.shape)


# In[ ]:


size=test_images.shape[0]
predict_labels=np.zeros(size)
errorCount = 0.0
train_images=np.mat(train_images)
test_images=np.mat(test_images)
for i in range(size):
    classNumStr=test_labels[i]
    Result = classify(test_images[i],train_images,train_labels,11)
    predict_labels[i]=Result
    #print("识别结果：%d 正确答案：%d" % (Result, classNumStr))
    if i %100 == 0:
        print("now is ",i)
        print("now error is \d",errorCount)
    if (Result != classNumStr):
        errorCount += 1.0
        #print("wrong num is %d",i)
print("\n错误数： %d" % errorCount)
fw = open("test.txt", 'w')    #将要输出保存的文件地址
for i in range(test_labels.shape[0]):    #读取的文件
    fw.write(str(test_labels[i])+ " " +str(predict_labels[i])) # 将字符串写入文件中
    fw.write("\n")
print("程序执行结束")


# In[ ]:


TP=np.zeros(10);FN=np.zeros(10);FP=np.zeros(10)
for i in range(test_labels.shape[0]):
    if test_labels[i]==predict_labels[i]:
        TP[int(test_labels[i])]+=1;
    else :
        FN[int(test_labels[i])]+=1
        FP[int(predict_labels[i])]+=1


# In[ ]:


ans=open("ans.txt", 'w')
for i in range(10):
    ans.write("Num is:"+str(i))
    ans.write(" Precision is "+str(TP[i]/(TP[i]+FP[i]))[0:5])
    ans.write(" Recall is "+str(TP[i]/(TP[i]+FN[i]))[0:5])
    ans.write("\n")

